# -*- coding: utf-8 -*-
"""
AI对话服务模块
实现与火山引擎API的集成和流式对话功能

主要功能:
- AI对话会话管理
- 流式响应处理
- 消息历史记录
- 上下文管理
- 错误处理和重试
"""

import json
import uuid
from typing import Dict, Any, List, Optional, Tuple, AsyncGenerator
from datetime import datetime, timezone
import logging
from dataclasses import dataclass, asdict
from enum import Enum

from .template_service import TemplateService
from .volcengine_service import get_volcengine_service, ChatMessage

# 配置日志
logger = logging.getLogger(__name__)


class MessageRole(Enum):
    """消息角色枚举"""
    USER = "user"
    ASSISTANT = "assistant"
    SYSTEM = "system"


class MessageType(Enum):
    """消息类型枚举"""
    TEXT = "text"
    IMAGE = "image"
    FILE = "file"
    SYSTEM = "system"
    ERROR = "error"


class SessionStatus(Enum):
    """会话状态枚举"""
    ACTIVE = "active"
    INACTIVE = "inactive"
    ARCHIVED = "archived"


@dataclass
class ChatContext:
    """对话上下文数据类"""
    session_id: str
    user_id: str
    requirement_id: Optional[str] = None
    template_id: Optional[str] = None
    template_variables: Optional[Dict[str, Any]] = None
    system_prompt: Optional[str] = None
    max_history: int = 10  # 最大历史消息数量
    temperature: float = 0.7
    max_tokens: int = 2000


@dataclass
class StreamResponse:
    """流式响应数据类"""
    session_id: str
    message_id: str
    content: str
    is_complete: bool = False
    tokens: Optional[Dict[str, int]] = None
    cost: Optional[float] = None
    response_time: Optional[float] = None
    error: Optional[str] = None


class ChatService:
    """AI对话服务类"""
    
    def __init__(self):
        self.template_service = TemplateService()
        
        # 内存存储
        self._session_storage = {}  # 存储会话数据
        self._message_storage = {}  # 存储消息数据
        self._session_counter = 0
        self._message_counter = 0
        
        # 对话配置
        self.default_system_prompt = """
你是一个专业的AI营销顾问，擅长制定营销策略和方案。
请根据用户的需求，提供专业、实用的营销建议和方案。
回答要条理清晰、逻辑性强，并提供具体可执行的建议。
"""
        
        # 初始化默认数据
        self._init_default_sessions()
    
    def _init_default_sessions(self):
        """初始化默认会话数据"""
        # 创建一个示例会话
        session_id = "session_001"
        self._session_storage[session_id] = {
            '_id': session_id,
            'user_id': 'user_001',
            'title': '营销方案咨询',
            'status': SessionStatus.ACTIVE.value,
            'requirement_id': None,
            'template_id': None,
            'template_variables': {},
            'system_prompt': self.default_system_prompt,
            'max_history': 10,
            'temperature': 0.7,
            'max_tokens': 2000,
            'message_count': 0,
            'total_tokens': 0,
            'total_cost': 0.0,
            'created_at': datetime.now(timezone.utc),
            'updated_at': datetime.now(timezone.utc)
        }
        self._session_counter = 1
    
    # 移除了会话创建功能，符合文档要求的简单设计
    
    def generate_marketing_plan(
        self,
        requirement_id: str,
        session_id: str,
        user_message: str = ""
    ) -> Dict[str, Any]:
        """
        生成营销方案
        
        Args:
            requirement_id: 需求ID
            session_id: 会话ID
            user_message: 用户消息（可选）
            
        Returns:
            Dict[str, Any]: 生成结果
        """
        try:
            # 获取或创建会话
            session = self._session_storage.get(session_id)
            if not session:
                # 创建新会话
                session = {
                    '_id': session_id,
                    'user_id': 'user_001',
                    'title': '营销方案生成',
                    'status': SessionStatus.ACTIVE.value,
                    'requirement_id': requirement_id,
                    'template_id': None,
                    'template_variables': {},
                    'system_prompt': self.default_system_prompt,
                    'max_history': 10,
                    'temperature': 0.7,
                    'max_tokens': 2000,
                    'message_count': 0,
                    'total_tokens': 0,
                    'total_cost': 0.0,
                    'created_at': datetime.now(timezone.utc),
                    'updated_at': datetime.now(timezone.utc)
                }
                self._session_storage[session_id] = session
            
            # 生成营销方案内容
            if user_message:
                content = self._generate_mock_ai_response(user_message)
            else:
                content = self._generate_default_marketing_plan(requirement_id)
            
            # 创建响应消息记录
            self._message_counter += 1
            response_id = f"msg_{self._message_counter:06d}"
            
            message_data = {
                '_id': response_id,
                'session_id': session_id,
                'role': MessageRole.ASSISTANT.value,
                'message_type': 'text',
                'content': content,
                'metadata': {
                    'requirement_id': requirement_id,
                    'tokens': {'prompt': 200, 'completion': 300, 'total': 500},
                    'cost': 0.01,
                    'response_time': 2.5
                },
                'created_at': datetime.now(timezone.utc)
            }
            
            self._message_storage[response_id] = message_data
            
            # 更新会话统计
            session['message_count'] += 1
            session['total_tokens'] += 500
            session['total_cost'] += 0.01
            session['updated_at'] = datetime.now(timezone.utc)
            
            return {
                'content': content,
                'session_id': session_id,
                'response_id': response_id
            }
            
        except Exception as e:
            logger.error(f"生成营销方案失败: {e}")
            raise ValueError(f"生成营销方案失败: {str(e)}")
    
    def _generate_default_marketing_plan(self, requirement_id: str) -> str:
        """
        生成默认营销方案
        
        Args:
            requirement_id: 需求ID
            
        Returns:
            str: 营销方案内容
        """
        return f"""
# 营销方案生成报告

## 方案概述
基于您提交的需求（ID: {requirement_id}），我们为您制定了以下营销方案：

## 核心策略
1. **目标定位**：精准锁定目标客户群体
2. **价值主张**：突出产品核心优势和差异化特点
3. **渠道策略**：多渠道整合营销，提升品牌曝光度

## 执行建议
### 第一阶段：市场调研与准备
- 深入了解目标客户需求和痛点
- 分析竞争对手策略和市场定位
- 制定详细的营销时间表

### 第二阶段：内容营销
- 创建高质量的内容素材
- 建立品牌故事和价值传播
- 优化SEO和社交媒体策略

### 第三阶段：推广执行
- 启动多渠道推广活动
- 监控数据表现和用户反馈
- 持续优化和调整策略

## 预期效果
- 提升品牌知名度30%
- 增加潜在客户转化率25%
- 实现销售目标增长20%

---
*本方案基于AI智能分析生成，建议结合实际情况进行调整优化。*
"""
    
    def send_message(
        self,
        session_id: str,
        user_message: str,
        message_type: str = "text",
        metadata: Optional[Dict[str, Any]] = None
    ) -> Tuple[bool, Optional[str], Optional[Dict[str, Any]]]:
        """
        发送消息并获取AI回复
        
        Args:
            session_id: 会话ID
            user_message: 用户消息
            message_type: 消息类型
            metadata: 消息元数据
            
        Returns:
            Tuple[bool, Optional[str], Optional[Dict]]: (成功标志, AI回复, 响应元数据)
        """
        try:
            # 获取会话信息
            session = self._session_storage.get(session_id)
            if not session:
                logger.error(f"会话不存在: {session_id}")
                return False, None, None
            
            # 创建用户消息记录
            self._message_counter += 1
            user_msg_id = f"msg_{self._message_counter:06d}"
            
            user_message_data = {
                '_id': user_msg_id,
                'session_id': session_id,
                'role': MessageRole.USER.value,
                'message_type': message_type,
                'content': user_message,
                'metadata': metadata or {},
                'created_at': datetime.now(timezone.utc)
            }
            
            self._message_storage[user_msg_id] = user_message_data
            
            # 模拟AI回复（实际项目中这里会调用真实的AI服务）
            ai_response = self._generate_mock_ai_response(user_message)
            
            # 创建AI回复消息记录
            self._message_counter += 1
            ai_msg_id = f"msg_{self._message_counter:06d}"
            
            ai_message_data = {
                '_id': ai_msg_id,
                'session_id': session_id,
                'role': MessageRole.ASSISTANT.value,
                'message_type': 'text',
                'content': ai_response,
                'metadata': {
                    'tokens': {'prompt': 100, 'completion': 150, 'total': 250},
                    'cost': 0.005,
                    'response_time': 1.2
                },
                'created_at': datetime.now(timezone.utc)
            }
            
            self._message_storage[ai_msg_id] = ai_message_data
            
            # 更新会话统计
            session['message_count'] += 2
            session['total_tokens'] += 250
            session['total_cost'] += 0.005
            session['updated_at'] = datetime.now(timezone.utc)
            
            # 构建响应元数据
            response_metadata = {
                'tokens': {'prompt': 100, 'completion': 150, 'total': 250},
                'cost': 0.005,
                'response_time': 1.2,
                'model': 'mock-model'
            }
            
            return True, ai_response, response_metadata
            
        except Exception as e:
            logger.error(f"发送消息失败 {session_id}: {e}")
            return False, None, None
    
    def _generate_mock_ai_response(self, user_message: str) -> str:
        """生成模拟AI回复"""
        if "营销" in user_message or "推广" in user_message:
            return """基于您的需求，我为您制定以下营销方案：

## 1. 目标市场分析
- 明确目标客户群体
- 分析市场竞争环境
- 识别市场机会点

## 2. 营销策略建议
- 多渠道整合营销
- 内容营销策略
- 社交媒体推广
- 线上线下结合

## 3. 执行计划
- 制定详细时间表
- 分配预算资源
- 设定关键指标
- 建立监控机制

请告诉我更多具体需求，我可以为您提供更详细的方案。"""
        else:
            return f"感谢您的咨询。针对您提到的'{user_message}'，我建议从以下几个方面考虑：\n\n1. 深入分析现状\n2. 制定明确目标\n3. 选择合适策略\n4. 制定执行计划\n\n请提供更多详细信息，我将为您提供更精准的建议。"
    
    async def send_message_stream(
        self,
        session_id: str,
        content: str,
        message_type: str = MessageType.TEXT.value,
        metadata: Optional[Dict[str, Any]] = None
    ) -> AsyncGenerator[StreamResponse, None]:
        """
        发送消息（流式响应）
        
        Args:
            session_id: 会话ID
            content: 消息内容
            message_type: 消息类型
            metadata: 元数据
            
        Yields:
            StreamResponse: 流式响应数据
        """
        assistant_message_id = None
        accumulated_content = ""
        start_time = datetime.now(timezone.utc)
        
        try:
            # 获取会话信息
            session = await self.session_dao.get_session_by_id(session_id)
            if not session:
                yield StreamResponse(
                    session_id=session_id,
                    message_id="",
                    content="",
                    error="会话不存在"
                )
                return
            
            # 创建用户消息
            user_message_id = await self._create_user_message(
                session_id, content, message_type, metadata
            )
            if not user_message_id:
                yield StreamResponse(
                    session_id=session_id,
                    message_id="",
                    content="",
                    error="创建用户消息失败"
                )
                return
            
            # 构建对话上下文
            context = await self._build_chat_context(session)
            
            # 获取历史消息
            history_messages = await self._get_history_messages(session_id, context.max_history)
            
            # 预创建AI消息记录
            assistant_message_id = await self._create_assistant_message(
                session_id=session_id,
                content="",  # 初始为空
                tokens={},
                cost=0.0,
                response_time=0.0
            )
            
            # 调用AI流式服务
            volcengine_service = await self._get_volcengine_service()
            
            async for chunk in volcengine_service.chat_completion_stream(
                messages=history_messages,
                temperature=context.temperature,
                max_tokens=context.max_tokens
            ):
                # 处理火山引擎API的响应格式
                if 'choices' in chunk and chunk['choices']:
                    choice = chunk['choices'][0]
                    if 'delta' in choice and 'content' in choice['delta']:
                        content = choice['delta']['content']
                        if content:
                            accumulated_content += content
                            
                            yield StreamResponse(
                                session_id=session_id,
                                message_id=assistant_message_id,
                                content=content,
                                is_complete=False
                            )
                elif 'error' in chunk:
                    yield StreamResponse(
                        session_id=session_id,
                        message_id=assistant_message_id or "",
                        content="",
                        error=chunk['error']
                    )
                    return
            
            # 流式响应完成
            end_time = datetime.now(timezone.utc)
            response_time = (end_time - start_time).total_seconds()
            
            # 更新AI消息内容
            if assistant_message_id and accumulated_content:
                await self.message_dao.update_message(
                    assistant_message_id,
                    {
                        'content': accumulated_content,
                        'response_time': response_time,
                        'updated_at': end_time
                    }
                )
                
                # 更新会话统计
                await self._update_session_stats(session_id, {}, response_time)
            
            # 发送完成信号
            yield StreamResponse(
                session_id=session_id,
                message_id=assistant_message_id or "",
                content=accumulated_content,
                is_complete=True,
                response_time=response_time
            )
            
        except Exception as e:
            logger.error(f"流式发送消息失败 {session_id}: {e}")
            
            # 创建错误消息
            if assistant_message_id:
                await self.message_dao.update_message(
                    assistant_message_id,
                    {
                        'content': f"系统错误: {str(e)}",
                        'message_type': MessageType.ERROR.value,
                        'updated_at': datetime.now(timezone.utc)
                    }
                )
            
            yield StreamResponse(
                session_id=session_id,
                message_id=assistant_message_id or "",
                content="",
                error=f"系统错误: {str(e)}"
            )
    
    # 移除了历史查询功能，符合文档要求的简单设计
    
    # 移除了会话删除功能，符合文档要求的简单设计
    
    # 移除了反馈处理功能，符合文档要求的简单设计
    
    async def _get_volcengine_service(self):
        """获取火山引擎服务实例"""
        return await get_volcengine_service()
    
# 全局服务实例
chat_service = ChatService()

def get_chat_service() -> ChatService:
    """获取聊天服务实例"""
    return chat_service
    
    async def _build_chat_context(self, session: Dict[str, Any]) -> ChatContext:
        """构建对话上下文"""
        return ChatContext(
            session_id=session['_id'],
            user_id=session['user_id'],
            requirement_id=session.get('requirement_id'),
            template_id=session.get('template_id'),
            template_variables=session.get('template_variables', {}),
            system_prompt=session.get('system_prompt'),
            max_history=session.get('max_history', 10),
            temperature=session.get('temperature', 0.7),
            max_tokens=session.get('max_tokens', 2000)
        )
    
    async def _get_history_messages(
        self,
        session_id: str,
        max_history: int
    ) -> List[ChatMessage]:
        """获取历史消息用于AI调用"""
        try:
            # 获取对话历史（排除系统消息和错误消息）
            messages = await self.message_dao.get_conversation_history(
                session_id=session_id,
                limit=max_history * 2  # 用户和助手消息成对出现
            )
            
            # 转换为ChatMessage格式
            chat_messages = []
            for msg in messages:
                if msg['role'] in [MessageRole.USER.value, MessageRole.ASSISTANT.value]:
                    chat_messages.append(ChatMessage(
                        role=msg['role'],
                        content=msg['content']
                    ))
            
            return chat_messages
            
        except Exception as e:
            logger.error(f"获取历史消息失败 {session_id}: {e}")
            return []
    
    async def _create_user_message(
        self,
        session_id: str,
        content: str,
        message_type: str,
        metadata: Optional[Dict[str, Any]] = None
    ) -> Optional[str]:
        """创建用户消息"""
        try:
            message_data = {
                'session_id': session_id,
                'role': MessageRole.USER.value,
                'message_type': message_type,
                'content': content,
                'metadata': metadata or {}
            }
            
            return await self.message_dao.create_message(message_data)
            
        except Exception as e:
            logger.error(f"创建用户消息失败: {e}")
            return None
    
    async def _create_assistant_message(
        self,
        session_id: str,
        content: str,
        tokens: Dict[str, int],
        cost: float,
        response_time: float
    ) -> Optional[str]:
        """创建AI助手消息"""
        try:
            message_data = {
                'session_id': session_id,
                'role': MessageRole.ASSISTANT.value,
                'message_type': MessageType.TEXT.value,
                'content': content,
                'tokens': tokens,
                'cost': cost,
                'response_time': response_time
            }
            
            return await self.message_dao.create_message(message_data)
            
        except Exception as e:
            logger.error(f"创建AI助手消息失败: {e}")
            return None
    
    async def _create_system_message(
        self,
        session_id: str,
        content: str
    ) -> Optional[str]:
        """创建系统消息"""
        try:
            message_data = {
                'session_id': session_id,
                'role': MessageRole.SYSTEM.value,
                'message_type': MessageType.SYSTEM.value,
                'content': content
            }
            
            return await self.message_dao.create_message(message_data)
            
        except Exception as e:
            logger.error(f"创建系统消息失败: {e}")
            return None
    
    async def _create_error_message(
        self,
        session_id: str,
        error_content: str
    ) -> Optional[str]:
        """创建错误消息"""
        try:
            message_data = {
                'session_id': session_id,
                'role': MessageRole.ASSISTANT.value,
                'message_type': MessageType.ERROR.value,
                'content': error_content
            }
            
            return await self.message_dao.create_message(message_data)
            
        except Exception as e:
            logger.error(f"创建错误消息失败: {e}")
            return None
    
    async def _update_session_stats(
        self,
        session_id: str,
        usage: Dict[str, int],
        response_time: float
    ):
        """更新会话统计信息"""
        try:
            update_data = {
                'last_activity_at': datetime.now(timezone.utc)
            }
            
            if usage:
                # 增加token使用量
                total_tokens = usage.get('total_tokens', 0)
                if total_tokens > 0:
                    session = await self.session_dao.get_session_by_id(session_id)
                    if session:
                        update_data['total_tokens'] = session.get('total_tokens', 0) + total_tokens
                        update_data['total_cost'] = session.get('total_cost', 0.0) + self._calculate_cost(usage)
            
            # 增加消息数量
            session = await self.session_dao.get_session_by_id(session_id)
            if session:
                update_data['message_count'] = session.get('message_count', 0) + 1
            
            await self.session_dao.update_session(session_id, update_data)
            
        except Exception as e:
            logger.error(f"更新会话统计失败 {session_id}: {e}")
    
    def _calculate_cost(self, usage: Dict[str, int]) -> float:
        """计算API调用成本"""
        try:
            # 火山引擎定价（示例价格，需要根据实际情况调整）
            prompt_tokens = usage.get('prompt_tokens', 0)
            completion_tokens = usage.get('completion_tokens', 0)
            
            # 假设价格：输入token 0.001元/1000token，输出token 0.002元/1000token
            prompt_cost = (prompt_tokens / 1000) * 0.001
            completion_cost = (completion_tokens / 1000) * 0.002
            
            return prompt_cost + completion_cost
            
        except Exception as e:
            logger.error(f"计算成本失败: {e}")
            return 0.0


# 全局服务实例
chat_service = ChatService()


async def get_chat_service() -> ChatService:
    """
    获取对话服务实例
    
    Returns:
        ChatService: 服务实例
    """
    return chat_service


# 便捷函数
async def create_marketing_chat(
    user_id: str,
    title: str,
    requirement_id: Optional[str] = None,
    template_id: Optional[str] = None
) -> Optional[str]:
    """
    创建营销对话会话的便捷函数
    
    Args:
        user_id: 用户ID
        title: 会话标题
        requirement_id: 需求ID
        template_id: 模板ID
        
    Returns:
        Optional[str]: 会话ID
    """
    return await chat_service.create_chat_session(
        user_id=user_id,
        title=title,
        requirement_id=requirement_id,
        template_id=template_id
    )


async def send_marketing_message(
    session_id: str,
    content: str
) -> Optional[str]:
    """
    发送营销对话消息的便捷函数
    
    Args:
        session_id: 会话ID
        content: 消息内容
        
    Returns:
        Optional[str]: AI回复内容
    """
    return await chat_service.send_message(session_id, content)